import pandas as pd
import numpy as np
import math
from scipy.optimize import brentq

# --- 请在这里定义您的目标概率 ---
# 我在这里使用 0.85 作为示例。
TARGET_PROBABILITY = 0.875

def calculate_p(week, bmi, age):
    """为单个成员计算概率 'p'。"""
    # 这是您在描述中提供的计算公式
    p_val = 6.4496 + 0.0426 * week - 0.1218 * bmi - 0.0451 * age
    p_val = p_val / math.sqrt(1 + 1.7 * 1.7 * 0.129 * 0.129)
    return 1.0 / (1 + math.exp(-p_val))

def average_p_objective_function(week, cluster_data, target_p):
    """
    供求解器使用的目标函数。
    它会根据给定的 'week' 计算整个分组的平均 'p' 值，
    并返回该值与目标概率之间的差值。
    """
    # 使用列表推导式为分组内的每个人计算p值
    p_values = [calculate_p(week, row['基准BMI'], row['年龄']) for index, row in cluster_data.iterrows()]
    # 计算p值的平均值
    average_p = np.mean(p_values)
    # 返回平均值与目标的差值
    return average_p - target_p

def find_week_for_clusters(file_path, target_p):
    """
    加载已分组的数据，并为每个组找到最优的 week 值。
    """
    try:
        df = pd.read_csv(file_path)
    except FileNotFoundError:
        print(f"错误：文件 '{file_path}' 未找到。请确保它位于正确的目录中。")
        return

    # 获取所有唯一的组标签并排序
    clusters = sorted(df['cluster'].unique())
    results = {}

    print(f"正在为每个分组寻找最小 'week' 值，以使平均概率达到 {target_p}...")
    print("-" * 30)

    for cluster_id in clusters:
        # 获取当前分组的所有数据
        cluster_data = df[df['cluster'] == cluster_id]

        try:
            # 使用寻根算法 (brentq) 来寻找使目标函数为零的 week 值。
            # 我们将 cluster_data 和 target_p 作为额外参数传递给目标函数。
            # 'week' 的搜索范围设置为 -400 到 400，与您原始代码一致。
            optimal_week = brentq(average_p_objective_function, -400, 400, args=(cluster_data, target_p))
            results[cluster_id] = optimal_week
            print(f"分组 {cluster_id}: 找到 week = {optimal_week:.4f}")

        except ValueError:
            # 如果在给定的 'week' 范围内无法达到目标概率，则可能会发生此错误。
            print(f"分组 {cluster_id}: 无法找到 week。对于此分组的数据，目标概率可能过高或过低。")
            results[cluster_id] = None

    return results

if __name__ == '__main__':
    # 这是由您的第二个脚本创建的文件
    clustered_file = 'data_with_clusters.csv'
    week_results = find_week_for_clusters(clustered_file, TARGET_PROBABILITY)

    print(week_results)

